package com.mcp.weather;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.mcp.weather.config.Config;
import com.mcp.weather.cache.SimpleCache;
import com.mcp.weather.service.AmapWeatherService;
import com.mcp.weather.util.JsonRpcHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.function.Function;

/**
 * MCP 服务器核心类
 */
public class WeatherMcpServer {
    private static final Logger logger = LoggerFactory.getLogger(WeatherMcpServer.class);
    private static final ObjectMapper objectMapper = JsonRpcHandler.getObjectMapper();

    private final Config config;
    private final SimpleCache cache;
    private final AmapWeatherService weatherService;

    public WeatherMcpServer(Config config) {
        this.config = config;
        this.cache = new SimpleCache();
        this.weatherService = new AmapWeatherService(config, cache);
    }

    /**
     * 处理 JSON-RPC 请求
     */
    public String handleRequest(String requestJson) {
        return JsonRpcHandler.handleRequest(requestJson, method -> {
            try {
                switch (method) {
                    case "initialize":
                        return handleInitialize();
                    case "tools/list":
                        return handleListTools();
                    case "tools/call":
                        return handleCallTool(requestJson);
                    default:
                        return null;
                }
            } catch (Exception e) {
                logger.error("处理方法失败: {}", method, e);
                return null;
            }
        });
    }

    /**
     * 处理 initialize 请求
     */
    private JsonNode handleInitialize() {
        ObjectNode result = objectMapper.createObjectNode();
        
        // 设置协议版本（MCP 协议版本）
        result.put("protocolVersion", "2024-11-05");
        
        // 设置服务器信息
        ObjectNode serverInfo = objectMapper.createObjectNode();
        serverInfo.put("name", "mcp-weather-java");
        serverInfo.put("version", "1.0.0");
        result.set("serverInfo", serverInfo);
        
        // 设置能力
        ObjectNode capabilities = objectMapper.createObjectNode();
        // tools 能力：空对象表示支持工具功能
        ObjectNode tools = objectMapper.createObjectNode();
        capabilities.set("tools", tools);
        result.set("capabilities", capabilities);
        
        logger.info("Initialize response generated: protocolVersion={}, serverName={}, serverVersion={}", 
                result.get("protocolVersion").asText(), 
                result.get("serverInfo").get("name").asText(),
                result.get("serverInfo").get("version").asText());
        return result;
    }

    /**
     * 处理 tools/list 请求
     */
    private JsonNode handleListTools() {
        ObjectNode result = objectMapper.createObjectNode();
        ArrayNode tools = objectMapper.createArrayNode();
        
        ObjectNode tool = objectMapper.createObjectNode();
        tool.put("name", "weather.get");
        tool.put("description", "查询高德天气。支持中文城市名或 adcode，extensions=base|all");
        
        ObjectNode inputSchema = objectMapper.createObjectNode();
        inputSchema.put("type", "object");
        
        ObjectNode properties = objectMapper.createObjectNode();
        
        ObjectNode cityProp = objectMapper.createObjectNode();
        cityProp.put("type", "string");
        cityProp.put("description", "城市名称（中文）或 adcode，如 \"青岛\" 或 \"370200\"");
        properties.set("city", cityProp);
        
        ObjectNode extensionsProp = objectMapper.createObjectNode();
        extensionsProp.put("type", "string");
        ArrayNode enumValues = objectMapper.createArrayNode();
        enumValues.add("base");
        enumValues.add("all");
        extensionsProp.set("enum", enumValues);
        extensionsProp.put("default", "base");
        extensionsProp.put("description", "base=实时天气，all=预报");
        properties.set("extensions", extensionsProp);
        
        inputSchema.set("properties", properties);
        
        ArrayNode required = objectMapper.createArrayNode();
        required.add("city");
        inputSchema.set("required", required);
        
        tool.set("inputSchema", inputSchema);
        tools.add(tool);
        
        result.set("tools", tools);
        return result;
    }

    /**
     * 处理 tools/call 请求
     */
    private JsonNode handleCallTool(String requestJson) {
        try {
            JsonNode request = objectMapper.readTree(requestJson);
            JsonNode params = request.get("params");
            
            if (params == null || !params.has("name")) {
                return createToolErrorResponse("工具名称不能为空");
            }

            String toolName = params.get("name").asText();
            if (!"weather.get".equals(toolName)) {
                return createToolErrorResponse("Unknown tool");
            }

            JsonNode arguments = params.has("arguments") ? params.get("arguments") : null;
            String city = arguments != null && arguments.has("city") 
                    ? arguments.get("city").asText() : null;
            String extensions = arguments != null && arguments.has("extensions")
                    ? arguments.get("extensions").asText() : "base";

            // 调用天气服务
            com.mcp.weather.model.WeatherResponse response = weatherService.queryWeather(city, extensions);

            // 构建 MCP 响应
            ObjectNode result = objectMapper.createObjectNode();
            ArrayNode content = objectMapper.createArrayNode();
            ObjectNode textContent = objectMapper.createObjectNode();
            textContent.put("type", "text");
            
            // 检查是否有错误（错误时 data 为 null 且 raw 包含 ErrorResponse）
            if (response.getData() == null && response.getRaw() instanceof com.mcp.weather.model.WeatherResponse.ErrorResponse) {
                // 错误响应
                textContent.put("text", objectMapper.writeValueAsString(response.getRaw()));
            } else {
                // 成功响应
                textContent.put("text", objectMapper.writeValueAsString(response));
            }
            
            content.add(textContent);
            result.set("content", content);
            
            return result;

        } catch (Exception e) {
            logger.error("处理工具调用失败", e);
            return createToolErrorResponse("处理工具调用失败: " + e.getMessage());
        }
    }

    /**
     * 创建工具错误响应
     */
    private JsonNode createToolErrorResponse(String errorMessage) {
        ObjectNode result = objectMapper.createObjectNode();
        ArrayNode content = objectMapper.createArrayNode();
        ObjectNode textContent = objectMapper.createObjectNode();
        textContent.put("type", "text");
        
        ObjectNode errorObj = objectMapper.createObjectNode();
        errorObj.put("error", errorMessage);
        textContent.put("text", objectMapper.valueToTree(errorObj).toString());
        
        content.add(textContent);
        result.set("content", content);
        return result;
    }

    /**
     * 获取天气服务实例（用于 HTTP 服务器共享缓存）
     */
    public AmapWeatherService getWeatherService() {
        return weatherService;
    }
}

